#include "WindowService.h"
#include "RegeditHelper.h"
#include "Convert.h"
#include <windows.h>
#include <string>
#include <tlhelp32.h>
#include <iostream>
#include <WtsApi32.h>
#include <tchar.h>
#include <strsafe.h>
#include <codecvt>
#include <filesystem>
#include <fstream>
#include <lmcons.h>
#include <chrono>
#include <iomanip>
#include <iosfwd>
#pragma comment(lib, "WtsApi32.lib")

namespace
{
SERVICE_STATUS          gSvcStatus;
SERVICE_STATUS_HANDLE   gSvcStatusHandle;
HANDLE                  ghSvcStopEvent = NULL;
FILE *fp{nullptr};
const auto PICASSO = L"PicassoService.exe";
const auto PICASSO_KEY = L"{3546DE13-789F-4C6F-AE34-7606D0C6C09F}";
const auto VANGOGHSERVICE = L"VanGoghService.exe";
const auto VANGOGHSERVICE_KEY = L"{DAAC1DB2-23E4-4143-94D4-D1B87DB91264}";
const auto NGINX = L"nginx.exe";
std::string getAppDataSVisionPath()
{
    std::string userName;
    constexpr auto formatStr =
        "C:\\Users\\{}\\AppData\\Roaming\\SVision\\Log\\PicassoHost\\";
    //as LocalSystem's %APPDATA% is `%systemroot%\System32\config\systemprofile`
    //so use current user's info to combine current user %APPDATA%
    while (true) {
        DWORD sessionCount;
        PWTS_SESSION_INFO pSessionInfo;
        if (WTSEnumerateSessions(WTS_CURRENT_SERVER_HANDLE, 0, 1, &pSessionInfo,
                                 &sessionCount))
            std::cout << "WTSEnumerateSessions error" << std::endl;
        bool ok = false;
        for (DWORD i = 0; i < sessionCount; i++) {
            if (pSessionInfo[i].State != WTSActive)
                continue;
            DWORD dwSessionId = pSessionInfo[i].SessionId;

            LPSTR pBuffer = NULL;
            DWORD dwSize = 0;
            WTSQuerySessionInformationA(WTS_CURRENT_SERVER_HANDLE, dwSessionId,
                                        WTSUserName, &pBuffer, &dwSize);
            userName = std::string(pBuffer, pBuffer + dwSize - 1);
            LocalFree(pBuffer);
            ok = true;
            if (dwSize > 0)
                return std::format(formatStr, userName);
        }
        if (ok)
            break;
        Sleep(500);
    }

    userName.resize(UNLEN + 1, '\0');
    DWORD  userNameSize = UNLEN + 1;
    if (GetUserNameA((char *)userName.data(), &userNameSize)) {
        userName.resize(userNameSize - 1);
        return std::format(formatStr, userName);
    }
    char *buf = nullptr;
    size_t sz = 0;
    if (_dupenv_s(&buf, &sz, "APPDATA") == 0 && buf != nullptr) {
        printf("APPDATA = %s\n", buf);
        free(buf);
    }
    if (!buf)
        return "";
    return std::string(buf) + "\\SVision\\Log\\PicassoHost\\";
}
std::string dateFormat(std::string fmt= "%Y-%m-%d")
{
    auto now = std::chrono::system_clock::now();
    std::time_t now_time = std::chrono::system_clock::to_time_t(now);
    std::tm local_time;
    if (localtime_s(&local_time, &now_time) == 0) {
        std::ostringstream oss;
        oss << std::put_time(&local_time, fmt.c_str());
        return oss.str();
    }
    return "";
}

std::string date;

#define LOG(message) fprintf(fp,"%s | %s\n", dateFormat("%Y-%m-%d %H:%M:%S").c_str(),(message));

void createLog()
{
    std::string logPath = getAppDataSVisionPath();
    if (fp)
        fclose(fp);
    date = dateFormat();
    auto dir = logPath + date;
    if (!std::filesystem::exists(dir))
        std::filesystem::create_directories(dir);
    auto file = dir + "\\" + "log.txt";
    errno_t err;
    err = fopen_s(&fp, file.c_str(), "a+");
    if (err != 0) {
        LOG("Error opening file.");
        return;
    }
	LOG("PicassoHost start.");
}


}


bool IsProcessRunning(const std::wstring &processName)
{
    HANDLE hSnap = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
    if (hSnap == INVALID_HANDLE_VALUE)
        return false;

    PROCESSENTRY32 pe32;
    pe32.dwSize = sizeof(PROCESSENTRY32);

    if (!Process32First(hSnap, &pe32)) {
        CloseHandle(hSnap);
        return false;
    }

    do {
        if (wcscmp(pe32.szExeFile, processName.c_str()) == 0) {
            CloseHandle(hSnap);
            return true;
        }
    } while (Process32Next(hSnap, &pe32));

    CloseHandle(hSnap);
    return false;
}


void ImpersonateActiveUserAndRun(const std::wstring &path,
                                 const std::wstring &exeName)
{
    DWORD sessionCount;
    PWTS_SESSION_INFO pSessionInfo;

    if (!WTSEnumerateSessions(WTS_CURRENT_SERVER_HANDLE, 0, 1, &pSessionInfo,
                              &sessionCount))
        return;
    for (DWORD i = 0; i < sessionCount; i++) {
        if (pSessionInfo[i].State != WTSActive)
            continue;
        DWORD dwSessionId = pSessionInfo[i].SessionId;
        HANDLE hToken;

        if (!WTSQueryUserToken(dwSessionId, &hToken)) {
            LOG(std::format("error WTSQueryUserToken{0}", GetLastError()));
            continue;
        }
        PROCESS_INFORMATION pi;
        STARTUPINFO si = { sizeof(si) };
        si.lpDesktop = (LPWSTR)
                       L"winsta0\\default"; // Use "winsta0\\default" for the default desktop of the interactive window station.
        si.dwFlags = STARTF_USESHOWWINDOW;
        si.wShowWindow = SW_SHOW;
        if (!SetCurrentDirectory(path.c_str()))
            LOG(std::format("Failed SetCurrentDirectory. Error code:{}\n", GetLastError()));
        if (!CreateProcessAsUser(hToken, (path + exeName).c_str(), NULL, NULL, NULL,
                                 FALSE, 0, NULL, NULL, &si, &pi)) {
            LOG(std::format("CreateProcessAsUser  ({})\n", GetLastError()));
            continue;
        }
        CloseHandle(pi.hProcess);
        CloseHandle(pi.hThread);
        CloseHandle(hToken);
    }
    WTSFreeMemory(pSessionInfo);
}
std::string w2s(const std::wstring& wstr)
{
	if (wstr.empty())
		return {};

	std::string strTo;
	int nSizeNeeded = ::WideCharToMultiByte(
		CP_UTF8, 0, &wstr[0], (int)wstr.size(), NULL, 0, NULL, NULL);
	strTo.resize(nSizeNeeded, 0);
	::WideCharToMultiByte(CP_UTF8,
		0,
		&wstr[0],
		(int)wstr.size(),
		&strTo[0],
		nSizeNeeded,
		NULL,
		NULL);
	return strTo;
}
void CheckAndRun(const std::wstring &key, const std::wstring &processName)
{
    auto getRegPath = [&processName](const std::wstring & key,
    std::wstring & path) {
        if (!Regedit::read(HKEY_LOCAL_MACHINE,
                           std::format(
                               LR"(SOFTWARE\WOW6432Node\Microsoft\Windows\CurrentVersion\Uninstall\{}_is1)",
                               key).c_str(),
                           path) &&
            !Regedit::read(HKEY_LOCAL_MACHINE,
                           std::format(LR"(SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall\{}_is1)",
                                       key).c_str(),
                           path)) {
            LOG(w2s(std::format(L"can not find {}.", processName.c_str())).c_str());
            return false;
        }
        return true;
    };

    if (IsProcessRunning(processName))
        return;
    LOG(w2s(std::format(L"{} is not running. Starting it now...", processName.c_str())).c_str());
    std::wstring path;
    if (processName == NGINX) {
        if (!getRegPath(PICASSO_KEY, path))
            return;
        path.append(L"nginx-1.24.0\\");
        ImpersonateActiveUserAndRun(path, NGINX);
    } else {
        if (!getRegPath(key, path))
            return;
        ImpersonateActiveUserAndRun(path, processName.c_str());
    }
    LOG(w2s(std::format(L"{} open success.", processName.c_str())).c_str());
}


void DeletePreService()
{
    SC_HANDLE schSCManager = OpenSCManager(NULL, NULL, SC_MANAGER_CONNECT);
    if (schSCManager == NULL) {
        std::cerr << "OpenSCManager failed!" << std::endl;
        return;
    }

    SC_HANDLE schService = OpenService(schSCManager, SVCNAME,
                                       SERVICE_STOP | SERVICE_QUERY_STATUS | DELETE);
    if (schService == NULL) {
        CloseServiceHandle(schSCManager);
        std::cerr << "OpenService failed!" << std::endl;
        return;
    }

    SERVICE_STATUS ss;
    if (!ControlService(schService, SERVICE_CONTROL_STOP, &ss))
        std::cerr << "ControlService failed!" << std::endl;

    if (!DeleteService(schService)) {
        CloseServiceHandle(schService);
        CloseServiceHandle(schSCManager);
        std::cerr << "DeleteService failed!" << std::endl;
        return;
    }
    std::cout << "Service deleted successfully!" << std::endl;
    CloseServiceHandle(schService);
    CloseServiceHandle(schSCManager);

}
//
// Purpose:
//   Installs a service in the SCM database
//
// Parameters:
//   None
//
// Return value:
//   None
//
VOID SvcInstall()
{
    DeletePreService();

    SC_HANDLE schSCManager;
    SC_HANDLE schService;
    TCHAR szPath[MAX_PATH];

    if (!GetModuleFileName(NULL, szPath,
                           MAX_PATH)) { //if (!GetModuleFileName(_T(""), szPath, MAX_PATH))
        printf("Cannot install service (%d)\n", GetLastError());
        return;
    }

    schSCManager = OpenSCManager(
                       NULL,                    // local computer
                       NULL,                    // ServicesActive database
                       SC_MANAGER_ALL_ACCESS);  // full access rights

    if (schSCManager == NULL) {
        printf("OpenSCManager failed (%d)\n", GetLastError());
        return;
    }

    schService = CreateService(
                     schSCManager,              // SCM database
                     SVCNAME,                   // name of service
                     SVCNAME,                   // service name to display
                     SERVICE_ALL_ACCESS,        // desired access
                     SERVICE_WIN32_OWN_PROCESS, // service type
                     SERVICE_AUTO_START,        // start type
                     SERVICE_ERROR_NORMAL,      // error control type
                     szPath,                    // path to service's binary
                     NULL,                      // no load ordering group
                     NULL,                      // no tag identifier
                     NULL,                      // no dependencies
                     NULL,                      // LocalSystem account
                     NULL);                     // no password

    if (schService == NULL) {
        printf("CreateService failed (%d)\n", GetLastError());
        CloseServiceHandle(schSCManager);
        return;
    }
    printf("Service installed successfully\n");
    if (!StartService(schService, 0, NULL))
        printf("StartService failed (%d)\n", GetLastError());
    CloseServiceHandle(schService);
    CloseServiceHandle(schSCManager);
}

//
// Purpose:
//   Entry point for the service
//
// Parameters:
//   dwArgc - Number of arguments in the lpszArgv array
//   lpszArgv - Array of strings. The first string is the name of
//     the service and subsequent strings are passed by the process
//     that called the StartService function to start the service.
//
// Return value:
//   None.
//
VOID WINAPI SvcMain(DWORD dwArgc, LPTSTR *lpszArgv)
{
    // Register the handler function for the service
    gSvcStatusHandle = RegisterServiceCtrlHandler(
                           SVCNAME,
                           SvcCtrlHandler);

    if (!gSvcStatusHandle) {
        SvcReportEvent((LPTSTR)TEXT("RegisterServiceCtrlHandler"));
        return;
    }

    // These SERVICE_STATUS members remain as set here
    gSvcStatus.dwServiceType = SERVICE_WIN32_OWN_PROCESS;
    gSvcStatus.dwServiceSpecificExitCode = 0;

    // Report initial status to the SCM
    ReportSvcStatus(SERVICE_START_PENDING, NO_ERROR, 3000);

    // Perform service-specific initialization and work.
    SvcInit(dwArgc, lpszArgv);
}
DWORD WINAPI Worker(LPVOID lParam)
{
    createLog();
    while (true) {
        if (date != dateFormat())
            createLog();
        CheckAndRun(PICASSO_KEY, PICASSO);
        CheckAndRun(VANGOGHSERVICE_KEY, VANGOGHSERVICE);
        Sleep(500);
        CheckAndRun(PICASSO_KEY, NGINX);
        Sleep(5000);
    }
}
//
// Purpose:
//   The service code
//
// Parameters:
//   dwArgc - Number of arguments in the lpszArgv array
//   lpszArgv - Array of strings. The first string is the name of
//     the service and subsequent strings are passed by the process
//     that called the StartService function to start the service.
//
// Return value:
//   None
//
VOID SvcInit(DWORD dwArgc, LPTSTR *lpszArgv)
{
    ghSvcStopEvent = CreateEvent(
                         NULL,    // default security attributes
                         TRUE,    // manual reset event
                         FALSE,   // not signaled
                         NULL);   // no name

    if (ghSvcStopEvent == NULL) {
        ReportSvcStatus(SERVICE_STOPPED, NO_ERROR, 0);
        return;
    }

    // Report running status when initialization is complete.
    ReportSvcStatus(SERVICE_RUNNING, NO_ERROR, 0);

    CreateThread(NULL, 0, Worker, 0, 0, 0);
    while (1) {
        // Check whether to stop the service.
        WaitForSingleObject(ghSvcStopEvent, INFINITE);
        ReportSvcStatus(SERVICE_STOPPED, NO_ERROR, 0);
        LOG("PicassoHost exit.");
        fclose(fp);
        return;
    }
}

VOID ReportSvcStatus(DWORD dwCurrentState,
                     DWORD dwWin32ExitCode,
                     DWORD dwWaitHint)
{
    static DWORD dwCheckPoint = 1;

    // Fill in the SERVICE_STATUS structure.

    gSvcStatus.dwCurrentState = dwCurrentState;
    gSvcStatus.dwWin32ExitCode = dwWin32ExitCode;
    gSvcStatus.dwWaitHint = dwWaitHint;

    if (dwCurrentState == SERVICE_START_PENDING)
        gSvcStatus.dwControlsAccepted = 0;
    else gSvcStatus.dwControlsAccepted = SERVICE_ACCEPT_STOP;

    if (dwCurrentState == SERVICE_RUNNING ||
        dwCurrentState == SERVICE_STOPPED)
        gSvcStatus.dwCheckPoint = 0;
    else gSvcStatus.dwCheckPoint = dwCheckPoint++;

    // Report the status of the service to the SCM.
    SetServiceStatus(gSvcStatusHandle, &gSvcStatus);
}

VOID WINAPI SvcCtrlHandler(DWORD dwCtrl)
{
    // Handle the requested control code.

    switch (dwCtrl) {
    case SERVICE_CONTROL_STOP:
        ReportSvcStatus(SERVICE_STOP_PENDING, NO_ERROR, 0);
        // Signal the service to stop.
        SetEvent(ghSvcStopEvent);
        ReportSvcStatus(gSvcStatus.dwCurrentState, NO_ERROR, 0);
    case SERVICE_CONTROL_INTERROGATE:
        break;
    default:
        break;
    }

}

//
// Purpose:
//   Logs messages to the event log
//
// Parameters:
//   szFunction - name of function that failed
//
// Return value:
//   None
//
// Remarks:
//   The service must have an entry in the Application event log.
//
VOID SvcReportEvent(LPTSTR szFunction)
{
    HANDLE hEventSource;
    LPCTSTR lpszStrings[2];
    TCHAR Buffer[80];

    hEventSource = RegisterEventSource(NULL, SVCNAME);

    if (NULL != hEventSource) {
        StringCchPrintf(Buffer, 80, TEXT("%s failed with %d"), szFunction,
                        GetLastError());

        lpszStrings[0] = SVCNAME;
        lpszStrings[1] = Buffer;

        ReportEvent(hEventSource,        // event log handle
                    EVENTLOG_ERROR_TYPE, // event type
                    0,                   // event category
                    0x00010,           // event identifier  SVC_ERROR
                    NULL,                // no security identifier
                    2,                   // size of lpszStrings array
                    0,                   // no binary data
                    lpszStrings,         // array of strings
                    NULL);               // no binary data

        DeregisterEventSource(hEventSource);
    }
}